/*
- C++ 实现的简单的CNN
- 依赖Opencv3 和 Eigen3
- 需要利用pytorch 上导出模型数据（存在.json 文件里）
- 实现的功能有
    - Conv2d
	- Relu
	- Max_pool2d
	- FLatten
	- Linear

- 使用方法
    - 加载模型
	    NNModule model(modelfile);
	- 使用模型 
	    model(input_pic);
	  和pytorch 类似
*/

#if(1)
#ifndef TORCH_H
#define TORCH_H

#include <Eigen/Dense>
#include <opencv2/core/eigen.hpp>
#include <vector>
#include <string>
#include <opencv2/core.hpp>
#include <memory>

class CpuTensor4d
{
public:
	Eigen::MatrixXf data;
	int shape[4];
public:
	CpuTensor4d(const std::vector<int>& shape);

	CpuTensor4d(CpuTensor4d&& t)noexcept;

	CpuTensor4d& operator=(CpuTensor4d&& t)noexcept;

	CpuTensor4d& operator=(const CpuTensor4d& t);

	CpuTensor4d(const CpuTensor4d& t);

	CpuTensor4d();

	Eigen::Map<Eigen::MatrixXf> operator()(int dim0, int dim1);
	Eigen::Map<const Eigen::MatrixXf> operator()(int dim0, int dim1)const;
	
	std::vector<int> getShape()const;
	friend class Conv2d;
	friend class Relu;
	friend class Max_pool2d;
	friend class Flatten;
	friend class Flatten;
	friend class Linear;
};
class GpuTensor4d
{

};

using Tensor4d = CpuTensor4d;

class NNLayer
{
public:
	virtual Tensor4d operator()(const Tensor4d& input) const = 0;
};

class Conv2d :public NNLayer
{
public:
	Tensor4d conv_w;
	Tensor4d conv_b;
    Eigen::MatrixXf conv(const Eigen::MatrixXf& filter, const Eigen::MatrixXf& input) const;
public: 

	Conv2d(const Tensor4d& w, const Tensor4d& b);
	Tensor4d operator()(const Tensor4d& input) const;
};

class Relu :public NNLayer
{
private:
	Eigen::MatrixXf relu(const Eigen::MatrixXf& input) const;
public:
	Relu();
	Tensor4d operator()(const Tensor4d& input) const;
};

class Max_pool2d :public NNLayer
{
private:
	int kernel_size;
	//stride = kernel_size
public:
	Max_pool2d(const int kernel_size);
	Tensor4d operator()(const Tensor4d& input) const;
};

class Flatten :public NNLayer
{
public:
	Flatten();
	Tensor4d operator()(const Tensor4d& input) const;
};

class Linear :public NNLayer
{
private:
	const Tensor4d fc_w;
	const Tensor4d fc_b;
public:
	Linear(const Tensor4d& w, const Tensor4d& b);
	Tensor4d operator()(const Tensor4d& input) const;
};


// NN模型
class NNModule
{
private:
	int width;
	int height;
	std::vector<NNLayer*>layers;
	Tensor4d calculate(const Tensor4d& input) const;
	void normalize(const cv::Mat& src, cv::Mat& dst)const;
public:
	bool load(const std::string& modelfile);
	NNModule();
	NNModule(const std::string& modelfile);
	std::vector<int> operator()(const std::vector<cv::Mat>& images) const;
	int operator()(const cv::Mat& image) const;
	~NNModule();
};

#endif
#endif